from project_utils.image_io import *
from project_utils.filtering.core import *

from argparse import ArgumentParser
from numpy.fft import fft2, fftshift, ifft2, ifftshift

import matplotlib.pyplot as plt
import numpy as np
import os.path as osp


def zero_pad(row, col, kernel):
    """Padding the kernel to the specified size with 0

    `row`: padding to target row size

    `col`: padding to target column size

    `kernel`: input kernel
    """

    rows_to_add = row - kernel.shape[0]
    cols_to_add = col - kernel.shape[1]
    rows_to_pad_top = rows_to_add // 2
    rows_to_pad_bottom = rows_to_add - rows_to_pad_top
    cols_to_pad_left = cols_to_add // 2
    cols_to_pad_right = cols_to_add - cols_to_pad_left
    kernel = np.pad(
        kernel,
        ((rows_to_pad_top, rows_to_pad_bottom), (cols_to_pad_left, cols_to_pad_right)),
        mode="constant",
        constant_values=0,
    )
    return kernel


def convfilt_2d(
    array: np.ndarray,
    kernel: np.ndarray,
    size: str = "same" or "full",
    mode: str = "constant",
    value: int = 0,
):
    """Apply custom 2D convolution kernels to the input array,
    which ultilizes numpy broadcasting operation.

    `array`: a numpy ndarray in the shape of (h, w, ...).

    `kernel`: the convolution kernel in the shape of (kh, kw, ...), it's
    suggusted that the remaining demensions match those of `array`.

    `size`: size of the output result.

    `mode`: the boundary mode specify how the array will be padded.

    `value`: a constant value that fills the padding in constant mode.

    Note that `mode` and `value` will be forwarded to function `padding_2d`.
    """

    assert size in ["same", "full"] and kernel.ndim >= 2

    ah, aw = array.shape[:2]
    kh, kw = kernel.shape[:2]

    full_array = padding_2d(array, (kh, kw), mode, value)

    def convfilt_kernel(row_window: np.ndarray):
        flat_kerkel_shape = (kh * kw,) + kernel.shape[2:]
        row_kernel = kernel.reshape(*flat_kerkel_shape)[::-1]

        max_dim = max(row_window.ndim, row_kernel.ndim)

        if row_window.ndim != max_dim:
            axis = list(range(row_window.ndim, max_dim))
            row_window = np.expand_dims(row_window, axis)
        if row_kernel.ndim != max_dim:
            axis = list(range(row_kernel.ndim, max_dim))
            row_kernel = np.expand_dims(row_kernel, axis)

        return np.sum(row_window * row_kernel, axis=0)

    out_array = spatial_filtering(full_array, (kh, kw), convfilt_kernel)

    if size == "same":
        ah, aw = array.shape[:2]
        oh, ow = out_array.shape[:2]

        hs = (oh - ah) // 2
        ws = (ow - aw) // 2

        filt_array = out_array[hs : hs + ah, ws : ws + aw]

    if size == "full":
        filt_array = out_array

    return filt_array


def imfilter(
    array: np.ndarray,
    kernel: np.ndarray,
    filtering_mode: str = "corr",
    size: str = "same",
    mode: str = "constant",
    value: int = 0,
):
    """A python version of the matlab function of the same name,
    which computes convolution or correlation.

    `array`: a numpy ndarray in the shape of (h, w, ...).

    `kernel`: the convolution kernel in the shape of (kh, kw, ...), it's
    suggusted that the remaining demensions match those of `array`.

    `filtering_mode`: array or kernel arithmetic operations.

    `size`: size of the output result.

    `mode`: the boundary mode specify how the array will be padded.

    `value`: a constant value that fills the padding in constant mode.
    """

    assert filtering_mode in ["corr", "conv"]

    if filtering_mode == "corr":
        kernel_flip = kernel.reshape(kernel.size)
        kernel_flip = kernel_flip[::-1]
        kernel_flip = kernel_flip.reshape(kernel.shape)
        filt_array = convfilt_2d(array, kernel_flip, size, mode, value)

    if filtering_mode == "conv":
        filt_array = convfilt_2d(array, kernel, size, mode, value)

    return filt_array


def fspecial(type, parameters=None, sigma=None):
    """A python version of the matlab function of the same name, which supports a variety
    of predefined 2D linear spatial filters and generates filter masks.

    `type`: selection of a specific 2D linear filtering method.

    `parameters`: input parameters which related to filtering method.

    `sigma`: input additional parameter if essential.
    """

    assert type in [
        "motion",
    ], "Please choose any of the above type of kernel."

    if type == "motion":
        parameters = parameters or 9
        sigma = sigma or 0
        assert parameters > 0, "Parameters should be a positive number."

        len = max(1, parameters)
        half = (len - 1) / 2  # rotate half length around center
        phi = (sigma % 180) / 180 * np.pi

        cosphi = np.cos(phi)
        sinphi = np.sin(phi)
        xsign = np.sign(cosphi)
        linewdt = 1

        # Define mesh for the half matrix, eps takes care of the right size for 0 & 90 rotation
        sx = int(half * cosphi + linewdt * xsign - len * np.finfo(float).eps)
        sy = int(half * sinphi + linewdt - len * np.finfo(float).eps)
        x = np.arange(0, sx + xsign, xsign)
        y = np.arange(0, sy + 1)
        x, y = np.meshgrid(x, y)

        # Define shortest distance from a pixel to the rotated line
        dist2line = y * cosphi - x * sinphi  # distance perpendicular to the line

        rad = np.sqrt(x**2 + y**2)
        # Find points beyond the line's end-point but within the line width
        lastpix = np.where((rad >= half) & (np.abs(dist2line) <= linewdt))

        # Distance to the line's end-point parallel to the line
        x2lastpix = half - np.abs((x[lastpix] + dist2line[lastpix] * sinphi) / cosphi)

        dist2line[lastpix] = np.sqrt(dist2line[lastpix] ** 2 + x2lastpix**2)
        dist2line = linewdt + np.finfo(float).eps - np.abs(dist2line)
        dist2line[dist2line < 0] = 0  # Zero out anything beyond the line width

        # Unfold half-matrix to the full size
        h = np.rot90(dist2line, 2)
        kernel = np.zeros((2 * h.shape[0] - 1, 2 * h.shape[1] - 1))
        kernel[0 : h.shape[0], 0 : h.shape[1]] = h
        kernel[
            h.shape[0] - 1 : 2 * h.shape[0], h.shape[1] - 1 : 2 * h.shape[1]
        ] = dist2line
        kernel = kernel / (np.sum(kernel) + np.finfo(float).eps * len**2)

        if cosphi > 0:
            kernel = np.flipud(kernel)

    return kernel


def direct_inverse_filtering(image, kernel):
    """Direct inverse filtering of the input image

    `image`: array of input image

    `kernel`: array of input kernel
    """

    epsilon = 0.05
    kernel = zero_pad(image.shape[0], image.shape[1], kernel)

    G = fftshift(fft2(image))
    H = fftshift(fft2(kernel))
    F = G / (H + epsilon)

    filtered_image = np.real(ifft2(ifftshift(F)))
    filtered_image = np.rot90(filtered_image, k=2)

    return filtered_image


def constrained_inverse_filtering(image, kernel, radius):
    """Frequency-limited inverse filtering of input images

    `image`: array of input image

    `kernel`: array of input kernel

    `radius`: cut-off radius
    """

    kernel = zero_pad(image.shape[0], image.shape[1], kernel)

    G = fftshift(fft2(image))
    H = fftshift(fft2(kernel))

    center = np.array(H.shape) // 2
    y, x = np.ogrid[: H.shape[0], : H.shape[1]]
    mask = (y - center[0]) ** 2 + (x - center[1]) ** 2 > radius**2

    H[mask] = 1
    F = G / H

    filtered_image = np.real(ifft2(ifftshift(F)))
    filtered_image = np.rot90(filtered_image, k=2)

    return filtered_image


def wiener_filtering(image, kernel, const):
    """Frequency-limited inverse filtering of input images

    `image`: array of input image

    `kernel`: array of input kernel

    `const`: input const
    """

    kernel = zero_pad(image.shape[0], image.shape[1], kernel)

    G = fftshift(fft2(image))
    H = fftshift(fft2(kernel))
    H_2 = np.abs(H) ** 2
    F = (H_2 * G) / (H * (H_2 + const))

    filtered_image = np.real(ifft2(ifftshift(F)))
    filtered_image = np.rot90(filtered_image, k=2)

    return filtered_image


def constrained_least_squares_filtering(image, kernel, gamma, step, alpha, mu, sigma):
    """Frequency-limited inverse filtering of input images

    `image`: array of input image

    `kernel`: array of input kernel

    `gamma`:  coefficients of the Laplace kernel

    `step`:  step size of the dichotomy

    `alpha`:  Thresholds for constraints

    `mu`: mean value of Gaussian noise.

    `sigma`: standard deviation of Gaussian noise.
    """

    flag = 0
    eta_2 = image.size * (mu**2 + sigma**2)
    p = np.array([[0, -1, 0], [-1, 4, -1], [0, -1, 0]])
    p = zero_pad(image.shape[0], image.shape[1], p)
    kernel = zero_pad(image.shape[0], image.shape[1], kernel)

    G = fftshift(fft2(image / 255.0))
    H = fftshift(fft2(kernel))
    P = fftshift(fft2(p))
    H_2 = np.abs(H) ** 2
    P_2 = np.abs(P) ** 2

    while 1:
        F = (H_2 * G) / (H * (H_2 + gamma * P_2))
        R = G - H * F
        r = np.real(ifft2(ifftshift(R))) * 255.0
        r_2 = np.sum(r**2)
        if r_2 < eta_2 - alpha:
            if flag == -1:
                step /= 2
            gamma += step
            flag = 1
        elif r_2 > eta_2 + alpha:
            if flag == 1:
                step /= 2
            gamma -= step
            flag = -1
        else:
            break

    filtered_image = np.real(ifft2(ifftshift(F))) * 255.0
    filtered_image = np.rot90(filtered_image, k=2)

    return filtered_image


def add_gaussian_noise(source_image, mu=0, sigma=0.1):
    """Adding Gaussian noise to input images

    `source_path`: source image path

    `target_path`: storage directory for the target image

    `mu`: mean value of Gaussian noise. Defaults to 0.

    `sigma`: standard deviation of Gaussian noise. Defaults to 0.1.
    """

    noise = np.random.normal(mu, sigma, source_image.shape)
    target_image = source_image + noise

    return target_image


def plot_comparison(figsize, images, titles, display=True):
    """Plot comparison between the source image and the target images."""

    fig = plt.figure(figsize=figsize)
    n_plots = min(len(images), len(titles))

    for idx, (image, title) in enumerate(zip(images, titles)):
        plt.subplot(1, n_plots, idx + 1)

        plt.imshow(image, cmap="gray")
        plt.title(label=title)
        plt.xlabel("X"); plt.ylabel("Y")

    plt.tight_layout()

    if display: plt.show()


def fspecial_test(source_path, target_path, len, theta):
    """Test "motion" convolution kernels appearing in the fspecial function."""

    source_image = load_image_as_grayscale(source_path, "rgb")
    kernel = fspecial("motion", len, theta)
    target_image = imfilter(source_image, kernel, "corr", mode="circular")
    target_image = np.clip(target_image, 0, 255)
    target_image = target_image.astype(np.uint8)

    target_path = osp.join(target_path, "motion.jpg")
    save_grayscale_image(target_path, target_image, False)

    plot_comparison(
        figsize=(10, 5),
        images=[
            source_image,
            target_image,
        ],
        titles=[
            "origin",
            "motion",
        ],
    )


def direct_inverse_filtering_test(source_path, target_path, len, theta):
    """Test direct inverse filtering function."""

    source_image = load_image_as_grayscale(source_path, "rgb")
    kernel = fspecial("motion", len, theta)
    target_image = direct_inverse_filtering(source_image, kernel)
    target_image = np.clip(target_image, 0, 255)
    target_image = target_image.astype(np.uint8)

    target_path = osp.join(target_path, "direct_inverse_filtering_motion.jpg")
    save_grayscale_image(target_path, target_image, False)

    plot_comparison(
        figsize=(10, 5),
        images=[
            source_image,
            target_image,
        ],
        titles=[
            "motion",
            "direct inverse filtering",
        ],
    )


def add_gaussian_noise_test(source_path, target_path, mu, sigma):
    """Test add gaussian noise function."""

    source_image = load_image_as_grayscale(source_path, "rgb")
    target_image = add_gaussian_noise(source_image, mu, sigma)
    target_image = np.clip(target_image, 0, 255)
    target_image = target_image.astype(np.uint8)

    target_path = osp.join(target_path, f"gaussian_{mu}_{sigma}.jpg")
    save_grayscale_image(target_path, target_image, False)

    plot_comparison(
        figsize=(10, 5),
        images=[
            source_image,
            target_image,
        ],
        titles=[
            "motion",
            "motion and gaussian",
        ],
    )


def constrained_inverse_filtering_test(
    source_path, target_path, len, theta, mu, sigma, radius
):
    """Test and compare direct inverse filtering function to constrained inverse filtering function."""

    source_image = load_image_as_grayscale(source_path, "rgb")
    kernel = fspecial("motion", len, theta)

    target_image_0 = add_gaussian_noise(source_image, mu, sigma)
    target_image_0 = np.clip(target_image_0, 0, 255)
    target_image_0 = target_image_0.astype(np.uint8)

    target_path_0 = osp.join(target_path, f"gaussian_{mu}_{sigma}.jpg")
    save_grayscale_image(target_path_0, target_image_0, False)

    target_image_1 = direct_inverse_filtering(target_image_0, kernel)
    target_image_1 = np.clip(target_image_1, 0, 255)
    target_image_1 = target_image_1.astype(np.uint8)

    target_path_1 = osp.join(target_path, "direct_inverse_filtering_gaussian.jpg")
    save_grayscale_image(target_path_1, target_image_1, False)

    target_image_2 = constrained_inverse_filtering(target_image_0, kernel, radius)
    target_image_2 = np.clip(target_image_2, 0, 255)
    target_image_2 = target_image_2.astype(np.uint8)

    target_path_2 = osp.join(target_path, "constrained_inverse_filtering_gaussian.jpg")
    save_grayscale_image(target_path_2, target_image_2, False)

    plot_comparison(
        figsize=(15, 5),
        images=[
            target_image_0,
            target_image_1,
            target_image_2,
        ],
        titles=[
            "motion and gaussian",
            "direct inverse filtering",
            "constrained inverse filtering",
        ],
    )


def wiener_filtering_test(source_path, target_path, len, theta, mu, sigma, const):
    """Test wiener filtering function."""

    source_image = load_image_as_grayscale(source_path, "rgb")
    kernel = fspecial("motion", len, theta)

    target_image_0 = add_gaussian_noise(source_image, mu, sigma)
    target_image_0 = np.clip(target_image_0, 0, 255)
    target_image_0 = target_image_0.astype(np.uint8)

    target_path_0 = osp.join(target_path, f"gaussian_{mu}_{sigma}.jpg")
    save_grayscale_image(target_path_0, target_image_0, False)

    target_image_1 = wiener_filtering(target_image_0, kernel, const)
    target_image_1 = np.clip(target_image_1, 0, 255)
    target_image_1 = target_image_1.astype(np.uint8)

    target_path_1 = osp.join(target_path, "wiener_gaussian.jpg")
    save_grayscale_image(target_path_1, target_image_1, False)

    plot_comparison(
        figsize=(10, 5),
        images=[
            target_image_0,
            target_image_1,
        ],
        titles=[
            "motion and gaussian",
            "wiener",
        ],
    )


def constrained_least_squares_filtering_test(
    source_path, target_path, len, theta, mu, sigma, gamma, step, alpha
):
    """Test constrained least squares filtering function."""

    source_image = load_image_as_grayscale(source_path, "rgb")
    kernel = fspecial("motion", len, theta)

    target_image_0 = add_gaussian_noise(source_image, mu, sigma)
    target_image_0 = np.clip(target_image_0, 0, 255)
    target_image_0 = target_image_0.astype(np.uint8)

    target_path_0 = osp.join(target_path, f"gaussian_{mu}_{sigma}.jpg")
    save_grayscale_image(target_path_0, target_image_0, False)

    target_image_1 = constrained_least_squares_filtering(
        target_image_0, kernel, gamma, step, alpha, mu, sigma
    )
    target_image_1 = np.clip(target_image_1, 0, 255)
    target_image_1 = target_image_1.astype(np.uint8)

    target_path_1 = osp.join(target_path, "constrained_least_squares_gaussian.jpg")
    save_grayscale_image(target_path_1, target_image_1, False)

    plot_comparison(
        figsize=(10, 5),
        images=[
            target_image_0,
            target_image_1,
        ],
        titles=[
            "motion and gaussian",
            "constrained least quares filtering",
        ],
    )


def main_procedure(cmdargs):
    if cmdargs.mode == "fspecial":
        source_path = osp.abspath(cmdargs.source)
        target_path = osp.abspath(cmdargs.target)
        fspecial_test(source_path, target_path, cmdargs.len, cmdargs.theta)

    elif cmdargs.mode == "direct_inverse":
        source_path = osp.abspath(cmdargs.source)
        target_path = osp.abspath(cmdargs.target)
        direct_inverse_filtering_test(
            source_path, target_path, cmdargs.len, cmdargs.theta
        )

    elif cmdargs.mode == "gaussian":
        source_path = osp.abspath(cmdargs.source)
        target_path = osp.abspath(cmdargs.target)
        add_gaussian_noise_test(source_path, target_path, cmdargs.mu, cmdargs.sigma)

    elif cmdargs.mode == "constrained_inverse":
        source_path = osp.abspath(cmdargs.source)
        target_path = osp.abspath(cmdargs.target)
        constrained_inverse_filtering_test(
            source_path,
            target_path,
            cmdargs.len,
            cmdargs.theta,
            cmdargs.mu,
            cmdargs.sigma,
            cmdargs.radius,
        )

    elif cmdargs.mode == "wiener":
        source_path = osp.abspath(cmdargs.source)
        target_path = osp.abspath(cmdargs.target)
        wiener_filtering_test(
            source_path,
            target_path,
            cmdargs.len,
            cmdargs.theta,
            cmdargs.mu,
            cmdargs.sigma,
            cmdargs.const,
        )

    elif cmdargs.mode == "constrained_least_squares":
        source_path = osp.abspath(cmdargs.source)
        target_path = osp.abspath(cmdargs.target)
        constrained_least_squares_filtering_test(
            source_path,
            target_path,
            cmdargs.len,
            cmdargs.theta,
            cmdargs.mu,
            cmdargs.sigma,
            cmdargs.gamma,
            cmdargs.step,
            cmdargs.alpha,
        )


if __name__ == "__main__":
    parser = ArgumentParser(description="Spatial Filtering Algorithms.")

    parser.add_argument(
        "mode",
        metavar="MODE",
        choices=[
            "fspecial",
            "direct_inverse",
            "gaussian",
            "constrained_inverse",
            "wiener",
            "constrained_least_squares",
        ],
        help="Select unit test: 'fspecial', 'direct_inverse', 'gaussian', 'constrained_inverse', 'wiener', 'constrained_least_squares'.",
    )

    parser.add_argument(
        "--source", metavar="SOURCE", type=str, help="Filepath of the source image."
    )
    parser.add_argument(
        "--target", metavar="TARGET", type=str, help="Folder of the target outputs."
    )

    parser.add_argument(
        "--len",
        metavar="LEN",
        type=float,
        help="Parameter of function 'fspecial'.",
        default=20,
    )
    parser.add_argument(
        "--theta",
        metavar="THETA",
        type=float,
        help="Parameter of function 'fspecial'.",
        default=60,
    )

    parser.add_argument(
        "--mu",
        metavar="MU",
        type=float,
        help="Parameter of function 'gaussian'.",
        default=0,
    )
    parser.add_argument(
        "--sigma",
        metavar="SIGMA",
        type=float,
        help="Parameter of function 'gaussian'.",
        default=20,
    )

    parser.add_argument(
        "--radius",
        metavar="RADIUS",
        type=float,
        help="Parameter of function 'constrained_inverse'.",
        default=20,
    )

    parser.add_argument(
        "--const",
        metavar="CONST",
        type=float,
        help="Parameter of function 'wiener'.",
        default=0.007,
    )

    parser.add_argument(
        "--gamma",
        metavar="GAMMA",
        type=float,
        help="Parameter of function 'constrained_least_squares'.",
        default=1,
    )
    parser.add_argument(
        "--step",
        metavar="STEP",
        type=float,
        help="Parameter of function 'constrained_least_squares'.",
        default=1,
    )
    parser.add_argument(
        "--alpha",
        metavar="ALPHA",
        type=float,
        help="Parameter of function 'constrained_least_squares'.",
        default=1,
    )

    main_procedure(parser.parse_args())
